
/****************************************************************
* @file   : pcie_tools.hpp
* @author : NingJian (freegodly@gmail.com)
* @date   : 05/26/21 09:37:48
* @brief   : 
****************************************************************/

#pragma once

#include "vpcie.hpp"
#include <exception>

#include <memory>
#include "eth_types.h"
#include <iostream>
#include "eth_types.h"
#include "eth_helper.h"
#include "eth_func.hpp"


using namespace std;
using namespace eth::mcore;


#define ACCESSES 64

static std::shared_ptr<eth_context> context;

__PCIE_MEM__ static uint64_t m_allocated_memory_dag = 0; 
__PCIE_MEM__ static size_t m_allocated_memory_light_cache = 0;


__PCIE_MEM__ static hash1024 * d_dag;
__PCIE_MEM__ static hash512 * d_light;

__PCIE_MEM__ static uint32_t d_dag_size;
__PCIE_MEM__ static uint32_t d_light_size;

__PCIE_MEM__ static hash256 d_header;
__PCIE_MEM__ static uint64_t d_target;

struct search_results
{
   bool    isfound{false};
   uint64_t nonce = 0;
   hash256 header_hash;
   hash256 result_hash;

};


bool pcie_initDevice()
{
    try
    {
        pcieSetDevice(1);
        pcieDeviceReset();
    }
    catch (std::exception & ec)
    {
        return false;
    }
    return true;
}


void pcie_get_constants(hash1024** _dag, uint32_t* _dag_size, hash512** _light, uint32_t* _light_size)
{
    if (_dag)
    {
        hash1024* _d;
        pcieMemcpyFromSymbol(&_d, d_dag, sizeof(hash1024*),0);
        *_dag = _d;
    }
    if (_dag_size)
    {
        uint32_t _ds;
        pcieMemcpyFromSymbol(&_ds, d_dag_size, sizeof(uint32_t),0);
        *_dag_size = _ds;
    }
    if (_light)
    {
        hash512* _l;
        pcieMemcpyFromSymbol(&_l, d_light, sizeof(hash512*),0);
        *_light = _l;
    }
    if (_light_size)
    {
        uint32_t _ls;
        pcieMemcpyFromSymbol(&_ls, d_light_size, sizeof(uint32_t),0);
        *_light_size = _ls;
    }
}

void pcie_set_constants(hash1024* _dag, uint32_t _dag_size, hash512* _light, uint32_t _light_size)
{
    pcieMemcpyToSymbol(d_dag, &_dag, sizeof(hash1024*),0);
    pcieMemcpyToSymbol(d_dag_size, &_dag_size, sizeof(uint32_t),0);
    pcieMemcpyToSymbol(d_light, &_light, sizeof(hash512*),0);
    pcieMemcpyToSymbol(d_light_size, &_light_size, sizeof(uint32_t),0);
}

void pcie_set_header(hash256 _header)
{
   pcieMemcpyToSymbol(d_header, &_header, sizeof(hash256),0 );
}

void pcie_set_target(uint64_t _target)
{
   pcieMemcpyToSymbol(d_target, &_target, sizeof(uint64_t),0 );
}

void pcie_ethash_calculate_dag_item(uint32_t need_node_index)
{
    union {
        hash1024  dag_node;
        uint64 dag_node_mem[25];
    };

    uint32_t node_index = need_node_index * 2;

    for(int i =0 ; i < 16; ++i)
    {
       dag_node.word32s[i] = d_light[node_index % d_light_size].word32s[i];
    }
    dag_node.word32s[0] ^= node_index;
    SHA3_512(dag_node_mem);
    for (int i = 0; i != 256; ++i)
    {
        uint32_t parent_index = fnv(node_index ^ i, dag_node.word32s[i % 16]) % d_light_size;
        for(int i =0 ; i < 16; ++i)
        {
           dag_node.word32s[i] = fnv(dag_node.word32s[i], d_light[parent_index].word32s[i]);
        }
    }
    SHA3_512(dag_node_mem);
    for(int i =0 ; i < 16; ++i)
    {
       d_dag[need_node_index].word32s[i] = dag_node.word32s[i];
    }

    

    node_index = need_node_index * 2 + 1;

    for(int i =0 ; i < 16; ++i)
    {
       dag_node.word32s[i] = d_light[node_index % d_light_size].word32s[i];
    }
    dag_node.word32s[0] ^= node_index;
    SHA3_512(dag_node_mem);
    for (int i = 0; i != 256; ++i)
    {
        uint32_t parent_index = fnv(node_index ^ i, dag_node.word32s[i % 16]) % d_light_size;
        for(int i =0 ; i < 16; ++i)
        {
           dag_node.word32s[i] = fnv(dag_node.word32s[i], d_light[parent_index].word32s[i]);
        }
    }
    SHA3_512(dag_node_mem);
    for(int i =0 ; i < 16; ++i)
    {
       d_dag[need_node_index].word32s[i+16] = dag_node.word32s[i];
    }
}

void pcie_ethash_generate_dag(uint64_t dag_size)
{
    for (uint32_t node_index = 0; node_index < dag_size; node_index++)
    {
        pcie_ethash_calculate_dag_item(node_index);
        pcieDeviceSynchronize();
    }
}



bool pcie_initEpoch_internal(std::shared_ptr<eth_context> ctx)
{
    context = ctx;
    bool retVar = false;
    auto startInit = std::chrono::steady_clock::now();
    size_t RequiredTotalMemory = ( context->light_cache_num_items * 64 + context->full_dataset_num_items * 128);
    size_t RequiredDagMemory = context->full_dataset_num_items * 128;

    try
    {
        hash1024* dag;
        hash512* light;

        // If we have already enough memory allocated, we just have to
        // copy light_cache and regenerate the DAG
        if (m_allocated_memory_dag < RequiredDagMemory ||
            m_allocated_memory_light_cache < context->light_cache_num_items * 64)
        {
           
            pcieDeviceReset();
            pcieSetDeviceFlags(0xFFFF);
            pcieDeviceSetCacheConfig({});

            
            pcieMalloc(reinterpret_cast<void**>(&light),context->light_cache_num_items * 64);

            m_allocated_memory_light_cache = context->light_cache_num_items * 64;
            pcieMalloc(reinterpret_cast<void**>(&dag),context->full_dataset_num_items * 128);

            m_allocated_memory_dag = context->full_dataset_num_items * 128;

           
        }
        else
        {
            pcie_get_constants(&dag, NULL, &light, NULL);
        }

        pcieMemcpy(reinterpret_cast<void*>(light), 
                   (void *)(context->light_cache),
                   context->light_cache_num_items * 64);

        pcie_set_constants(dag,context->full_dataset_num_items, light,
             context->light_cache_num_items); 

        //ethash_generate_dag(context->full_dataset_num_items);
        //pcie_ethash_generate_dag(100);
        pcieMemcpy(reinterpret_cast<void*>(dag), 
                   (void *)(context->full_dataset),
                   context->full_dataset_num_items * 128);

        
        retVar = true;
    }
    catch (std::exception & ec)
    {
        retVar = false;
    }

    return retVar;
}

void pcie_debug()
{
   cout << "d_dag[0] : " << hash1024_to_string(d_dag[0]) << endl;
   cout << "d_dag[1] : " << hash1024_to_string(d_dag[1]) << endl;
   cout << "d_dag[2] : " << hash1024_to_string(d_dag[2]) << endl;
}



bool pcie_compute_hash(uint64_t nonce, hash256& header, search_results & sr)
{
    uint64 state[12]={0};
    // state[0] = header.word64s[0];
    // state[1] = header.word64s[1];
    // state[2] = header.word64s[2];
    // state[3] = header.word64s[3];
    state[0].x = header.word32s[0];
    state[0].y = header.word32s[1];
    state[1].x = header.word32s[2];
    state[1].y = header.word32s[3];
    state[2].x = header.word32s[4];
    state[2].y = header.word32s[5];
    state[3].x = header.word32s[6];
    state[3].y = header.word32s[7];
    //state[4] = nonce;
    state[4].x = nonce & 0xffffffff ;
    state[4].y = nonce >> 32;
    keccak_f1600_init(state);

    uint128 mix[8] ={
        {state[0].x, state[0].y, state[1].x, state[1].y},
        {state[2].x, state[2].y, state[3].x, state[3].y},
        {state[4].x, state[4].y, state[5].x, state[5].y},
        {state[6].x, state[6].y, state[7].x, state[7].y},
        {state[0].x, state[0].y, state[1].x, state[1].y},
        {state[2].x, state[2].y, state[3].x, state[3].y},
        {state[4].x, state[4].y, state[5].x, state[5].y},
        {state[6].x, state[6].y, state[7].x, state[7].y},
    };

    unsigned init0 = state[0].x;

    unsigned int offset[ACCESSES];
    
    for(int i=0; i<ACCESSES; i++){   
        offset[i] = fnv(init0 ^ (i), ((unsigned int*)&mix[(i/4)%8])[i%4]) % d_dag_size;

        if( d_dag[offset[i]].word64s[0] == 0 ){
            pcie_ethash_calculate_dag_item(offset[i]);
        }

        mix[0].x = fnv(mix[0].x, d_dag[offset[i]].word32s[0] );
        mix[0].y = fnv(mix[0].y, d_dag[offset[i]].word32s[1] );
        mix[0].z = fnv(mix[0].z, d_dag[offset[i]].word32s[2] );
        mix[0].w = fnv(mix[0].w, d_dag[offset[i]].word32s[3] );
        mix[1].x = fnv(mix[1].x, d_dag[offset[i]].word32s[4] );
        mix[1].y = fnv(mix[1].y, d_dag[offset[i]].word32s[5] );
        mix[1].z = fnv(mix[1].z, d_dag[offset[i]].word32s[6] );
        mix[1].w = fnv(mix[1].w, d_dag[offset[i]].word32s[7] );
        mix[2].x = fnv(mix[2].x, d_dag[offset[i]].word32s[8] );
        mix[2].y = fnv(mix[2].y, d_dag[offset[i]].word32s[9]);
        mix[2].z = fnv(mix[2].z, d_dag[offset[i]].word32s[10]);
        mix[2].w = fnv(mix[2].w, d_dag[offset[i]].word32s[11]);
        mix[3].x = fnv(mix[3].x, d_dag[offset[i]].word32s[12]);
        mix[3].y = fnv(mix[3].y, d_dag[offset[i]].word32s[13]);
        mix[3].z = fnv(mix[3].z, d_dag[offset[i]].word32s[14]);
        mix[3].w = fnv(mix[3].w, d_dag[offset[i]].word32s[15]);
        mix[4].x = fnv(mix[4].x, d_dag[offset[i]].word32s[16]);
        mix[4].y = fnv(mix[4].y, d_dag[offset[i]].word32s[17]);
        mix[4].z = fnv(mix[4].z, d_dag[offset[i]].word32s[18]);
        mix[4].w = fnv(mix[4].w, d_dag[offset[i]].word32s[19]);
        mix[5].x = fnv(mix[5].x, d_dag[offset[i]].word32s[20]);
        mix[5].y = fnv(mix[5].y, d_dag[offset[i]].word32s[21]);
        mix[5].z = fnv(mix[5].z, d_dag[offset[i]].word32s[22]);
        mix[5].w = fnv(mix[5].w, d_dag[offset[i]].word32s[23]);
        mix[6].x = fnv(mix[6].x, d_dag[offset[i]].word32s[24]);
        mix[6].y = fnv(mix[6].y, d_dag[offset[i]].word32s[25]);
        mix[6].z = fnv(mix[6].z, d_dag[offset[i]].word32s[26]);
        mix[6].w = fnv(mix[6].w, d_dag[offset[i]].word32s[27]);
        mix[7].x = fnv(mix[7].x, d_dag[offset[i]].word32s[28]);
        mix[7].y = fnv(mix[7].y, d_dag[offset[i]].word32s[29]);
        mix[7].z = fnv(mix[7].z, d_dag[offset[i]].word32s[30]);
        mix[7].w = fnv(mix[7].w, d_dag[offset[i]].word32s[31]);
    }

    unsigned int thread_mix[8];
    for(int i=0; i<8; i++){
        thread_mix[i] = fnv_reduce(mix[i]);
    }

    state[8].x  = thread_mix[0]; state[8].y  = thread_mix[1];
    state[9].x  = thread_mix[2]; state[9].y  = thread_mix[3];
    state[10].x = thread_mix[4]; state[10].y = thread_mix[5];
    state[11].x = thread_mix[6]; state[11].y = thread_mix[7];
 

    uint64 calc_target = cuda_swab64(keccak_f1600_final(state));
    uint64_t r_calc_target = ((uint64_t)calc_target.y << 32 ) + calc_target.x;
    if (r_calc_target < d_target)
    {
        sr.isfound = true;
        sr.nonce = nonce;
        sr.header_hash = header;
        for(int i=0; i<8; i++){
            sr.result_hash.word32s[i] =  thread_mix[i];
        }

        return true;
    }
    return false;
}



bool pcie_ethash_search(search_results & sr, uint64_t start_nonce, uint64_t bach_size)
{
    for(uint64_t i = 0 ; i < bach_size; ++i) 
    {
        if(pcie_compute_hash(start_nonce + i, d_header, sr)) {
            return true;
        }

    }
    return false; 
}
